{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "a85aba8b",
   "metadata": {},
   "source": [
    "Dependencies"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "eb5e24d9",
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "# Install dependencies\n",
    "!pip install python-dotenv\n",
    "!pip install openai\n",
    "!pip install openpyxl"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "263b333f",
   "metadata": {},
   "outputs": [],
   "source": [
    "openai.api_key = ('') # Set api key"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c1c837c9",
   "metadata": {},
   "source": [
    "# No personality traits prompt"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c74b190a",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "\n",
    "import dotenv\n",
    "import json\n",
    "import re\n",
    "import openai\n",
    "import time\n",
    "import sqlite3\n",
    "import json\n",
    "\n",
    "from dotenv import load_dotenv\n",
    "load_dotenv()\n",
    "\n",
    "from threading import Thread\n",
    "\n",
    "MAX_FAILURES = 4\n",
    "\n",
    "\n",
    "def create_prompt(scenario, personality, author):\n",
    "    left_a, left_b = scenario[0]\n",
    "    right_a, right_b = scenario[1]\n",
    "    return f\"\"\"You are the fictional {personality}. \n",
    "- Your task is to make a decision in a game that determines a payment to you and Stranger A.\n",
    "- Stranger A has no choice in this game. You and Stranger A do not meet and will remain anonymous.\n",
    "- Make your decision using what you know of {personality} from the original {author} work and all subsequent commentary.\n",
    "- In this game you have two options:\n",
    "\n",
    "Option Left: Stranger A receives {left_a}, and you receive {left_b}, \n",
    "Option Right: Stranger A receives {right_a}, and you receive {right_b}, \n",
    "\n",
    "What do you choose, with one word [Left, Right]?\n",
    "\n",
    "Output three columns: 1) the choice you made [Left or Right], 2) the amount you receive, 3) and the amount Stranger A receives. Do not output anything else. An example of the format of the output is: Left, 750, 500\n",
    "\"\"\"\n",
    "# Should change example amount based on scenarios\n",
    "\n",
    "\n",
    "def get_decision(scenario, personality, author, scenario_name, model):\n",
    "    prompt = create_prompt(scenario, personality, author)\n",
    "    failure_count = 0\n",
    "    message_history = []\n",
    "    message_history.append({\"role\":\"user\", \"content\": prompt})\n",
    "    # Get response\n",
    "    completion = openai.ChatCompletion.create(\n",
    "        model='gpt-4',\n",
    "        messages=message_history,\n",
    "    )\n",
    "    reply_content1 = completion.choices[0].message.content\n",
    "\n",
    "    return dict({\"raw_reply\": reply_content1,\n",
    "                 \"choice\": \"Left\" if \"left\" in reply_content1.lower() else \"Right\", \n",
    "                 \"scenario\": scenario,\n",
    "                 \"personality\":personality, # Character\n",
    "                 \"model\":model,\n",
    "                 \"scenario_name\":scenario_name,\n",
    "                 \"prompt\":prompt\n",
    "                })"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "bd45f660",
   "metadata": {},
   "source": [
    "# Personality traits prompt"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8fd2e452",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "\n",
    "import dotenv\n",
    "import json\n",
    "import re\n",
    "import openai\n",
    "import time\n",
    "import sqlite3\n",
    "import json\n",
    "\n",
    "from dotenv import load_dotenv\n",
    "load_dotenv()\n",
    "\n",
    "from threading import Thread\n",
    "\n",
    "MAX_FAILURES = 4\n",
    "\n",
    "\n",
    "\n",
    "def create_prompt(scenario, personality, author):\n",
    "    left_a, left_b = scenario[0]\n",
    "    right_a, right_b = scenario[1]\n",
    "    return (f\"\"\"You are the fictional {personality}. \n",
    "- Your task is to make a decision in a game that determines a payment to you and Person A.\n",
    "- Person A has no choice in this game. You and Person A do not meet and will remain anonymous.\n",
    "- Make your decision using what you know of {personality} from the original {author} work and all subsequent commentary.\n",
    "- In this game you have two options.\"\"\",\n",
    "            \n",
    "f\"\"\"Option Left: Person A receives {left_a}, and you receive {left_b}. \n",
    "Option Right: Person A receives {right_a}, and you receive {right_b}.\n",
    "\n",
    "What personality traits of {personality} from {author}  would affect {personality}'s decision? Output a list of the 5 most impactful personality traits with little to no explanation:\"\"\",\n",
    "            \n",
    "f\"\"\"Based on these personality traits, what would {personality} choose, with one word [Left, Right]?\n",
    "Output two columns: 1) the choice you made [Left or Right], 2) the amount you receive. Do not output anything else. An example of the format of the output is: Right, 700\"\"\")\n",
    "\n",
    "\n",
    "\n",
    "def get_decision(scenario, personality, author, scenario_name, model):\n",
    "    prompt = create_prompt(scenario, personality, author)\n",
    "    failure_count = 0\n",
    "    message_history = []\n",
    "    message_history.append({\"role\":\"system\", \"content\": prompt[0]}) # System settings\n",
    "    message_history.append({\"role\":\"user\", \"content\": prompt[1]}) # Prompt 1\n",
    "    # Get response\n",
    "    completion = openai.ChatCompletion.create(\n",
    "        model='gpt-4',\n",
    "        messages=message_history,\n",
    "    )\n",
    "    reply_content1 = completion.choices[0].message.content\n",
    "    message_history.append({\"role\":\"assistant\", \"content\": reply_content1})\n",
    "    \n",
    "    message_history.append({\"role\":\"user\", \"content\": prompt[2]}) # Prompt 2\n",
    "    # Get response\n",
    "    completion = openai.ChatCompletion.create(\n",
    "        model='gpt-4',\n",
    "        messages=message_history,\n",
    "    )\n",
    "    reply_content2 = completion.choices[0].message.content\n",
    "    message_history.append({\"role\":\"assistant\", \"content\": reply_content2})\n",
    "\n",
    "    return dict({\"raw_reply\": message_history,\n",
    "                 \"choice\": \"Left\" if \"left\" in reply_content2.lower() else \"Right\", \n",
    "                 \"scenario\": scenario,\n",
    "                 \"personality\":personality, # Character\n",
    "                 \"model\":model,\n",
    "                 \"scenario_name\":scenario_name,\n",
    "                 \"prompt\":prompt[1],\n",
    "                 \"personalities\": reply_content1\n",
    "                })\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "629ebe4f",
   "metadata": {},
   "source": [
    "Games to run:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "643032f3",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Berk and Barc\n",
    "scenarios = dict({ \n",
    "    \"Berk29\": ((400, 400), (750, 400)),\n",
    "    \"Barc2\": ((400, 400), (750, 375)),\n",
    "    \"Berk23\": ((800, 200), (0, 0)),\n",
    "    \"Barc8\": ((300, 600), (700, 500)),\n",
    "    \"Berk15\": ((200, 700), (600, 600)),\n",
    "    \"Berk26\":((0, 800), (400, 400))\n",
    "})"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b6a4e0b9",
   "metadata": {},
   "source": [
    "Choose characters and run experiment"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5261014a",
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "from time import sleep, perf_counter\n",
    "import concurrent.futures\n",
    "\n",
    "class DictatorGame():\n",
    "\n",
    "    def __init__(self, scenario, personality, author, scenario_name, model):\n",
    "        self.scenario = scenario\n",
    "        self.personality = personality\n",
    "        self.scenario_name = scenario_name\n",
    "        self.model = model\n",
    "        self.author = author\n",
    "\n",
    "def run_game(args):\n",
    "    game = args[0]\n",
    "    err = 0\n",
    "    while err < 5:\n",
    "        try:\n",
    "            decision = get_decision(scenario = game.scenario, personality = game.personality, author = game.author, scenario_name = game.scenario_name, model = game.model)\n",
    "            return decision\n",
    "        except Exception as e:\n",
    "            err+=1\n",
    "            sleep(240)\n",
    "            print(f\"Error #{err}: {e}\")\n",
    "\n",
    "model = 'gpt-4'\n",
    "# List of personalities\n",
    "personalities = []\n",
    "\n",
    "# List of authors corresponding to personalities\n",
    "authors = []\n",
    "\n",
    "observations = list()\n",
    "\n",
    "with concurrent.futures.ThreadPoolExecutor() as executor:\n",
    "    for i, personality in enumerate(personalities):\n",
    "        for scenario_name, scenario in scenarios.items():\n",
    "            dictator_game = DictatorGame(scenario = scenario, personality = personality, author=authors[i], scenario_name = scenario_name, model = model)\n",
    "            f1 = executor.submit(run_game, args=[dictator_game])\n",
    "            observations.append(f1)\n",
    "                \n",
    "    # Check for completed results\n",
    "    i = 0\n",
    "    for f in concurrent.futures.as_completed(observations):\n",
    "        print(f'{i}')\n",
    "        i+=1"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "381872e9",
   "metadata": {},
   "source": [
    "# With personality traits"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "61207fd8",
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "import pandas as pd\n",
    "\n",
    "# Take observations and format the data into a pandas dataframe\n",
    "def write_data1(obv=observations):\n",
    "    df = pd.DataFrame(columns=['Model', 'Scenario', 'Character', 'Decision', 'RawData', 'Personalities'])\n",
    "    for i, observation in enumerate(obv):\n",
    "        result = observation.result()\n",
    "        new_row = pd.DataFrame({'Model': [result['model']], 'Scenario' : [result['scenario']], 'Character' : [result['personality']], \n",
    "                   'Decision' : [result['choice']], 'Personalities' : [result['personalities']]})\n",
    "        df = pd.concat([df, new_row], ignore_index=True)\n",
    "    return df\n",
    "df = write_data1(observations)    "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "cae7ff0d",
   "metadata": {},
   "source": [
    "# Without peronsality traits"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c48b4eeb",
   "metadata": {},
   "outputs": [],
   "source": [
    "import pandas as pd\n",
    "\n",
    "# Take observations and format the data into a pandas dataframe\n",
    "def write_data2(observations):\n",
    "    df = pd.DataFrame(columns=['Model', 'Scenario', 'Character', 'Decision', 'RawData'])\n",
    "    for observation in observations:\n",
    "        result = observation.result()\n",
    "        new_row = pd.DataFrame({'Model': [result['model']], 'Scenario' : [result['scenario']], 'Character' : [result['personality']], \n",
    "                        'Decision' : [result['choice']], 'RawData' : [result['raw_reply']]})\n",
    "        df = pd.concat([df, new_row], ignore_index=True)\n",
    "    return df\n",
    "df = write_data2(observations)    "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "bdc490b6",
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "# Write data to excel file\n",
    "df.to_excel('name.xlsx', index=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "20f30ead",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.9.7"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
